import requests
import json
import csv
import time
import re
from typing import List, Dict, Tuple

class RAGQAGenerator:
    """RAG问答对生成器"""
    
    def __init__(self, api_key: str, api_base: str = "https://api.openai.com/v1"):
        """
        初始化生成器
        
        Args:
            api_key: OpenAI API密钥
            api_base: API基础URL
        """
        self.api_key = api_key
        self.api_base = api_base
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
    
    def call_llm(self, content: str) -> str:
        """
        调用大语言模型生成问答对
        
        Args:
            content: 输入的法律条文内容
            
        Returns:
            模型返回的问答对文本
        """
        url = f"{self.api_base}/chat/completions"
        
        # 构建提示词
        prompt = f"""请根据以下宪法条文内容，生成3-5个高质量的问答对。问答对应该涵盖条文的核心要点、重要概念和实际应用场景。

要求：
1. 问题要具体、明确，便于用户理解
2. 答案要准确、完整，基于条文内容
3. 问题类型可以包括：概念解释、权利义务、制度规定、实际应用等
4. 每个问答对用"Q: "和"A: "标记

宪法条文内容：
{content}

请生成问答对："""

        data = {
            "model": "gpt-3.5-turbo",
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业的法律知识问答助手，擅长将法律条文转化为易于理解的问答形式。"
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            "temperature": 0.7,
            "max_tokens": 1000
        }
        
        try:
            response = requests.post(url, headers=self.headers, data=json.dumps(data))
            response.raise_for_status()
            result = response.json()
            return result['choices'][0]['message']['content']
        except Exception as e:
            print(f"API调用失败: {e}")
            return ""
    
    def parse_qa_pairs(self, llm_response: str) -> List[Tuple[str, str, List[str]]]:
        """
        解析LLM返回的问答对
        
        Args:
            llm_response: LLM返回的文本
            
        Returns:
            问答对列表，每个元素为(标题, 内容, 问题列表)
        """
        qa_pairs = []
        
        # 使用正则表达式匹配Q:和A:格式
        qa_pattern = r'Q:\s*(.*?)(?=\nA:|\nQ:|$)'
        a_pattern = r'A:\s*(.*?)(?=\nQ:|$)'
        
        questions = re.findall(qa_pattern, llm_response, re.DOTALL)
        answers = re.findall(a_pattern, llm_response, re.DOTALL)
        
        # 确保问题和答案数量匹配
        min_len = min(len(questions), len(answers))
        
        for i in range(min_len):
            question = questions[i].strip()
            answer = answers[i].strip()
            
            # 生成标题（取问题的前20个字符）
            title = question[:20] + "..." if len(question) > 20 else question
            
            # 将问题作为问题列表
            question_list = [question]
            
            qa_pairs.append((title, answer, question_list))
        
        return qa_pairs
    
    def read_split_result(self, file_path: str) -> List[Dict]:
        """
        读取拆分结果文件
        
        Args:
            file_path: 拆分结果文件路径
            
        Returns:
            拆分单元列表
        """
        units = []
        current_unit = {}
        
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        for line in lines:
            line = line.strip()
            if line.startswith('=== '):
                # 保存之前的单元
                if current_unit:
                    units.append(current_unit)
                
                # 解析新单元
                parts = line.strip('= ').split(' ')
                if len(parts) >= 2:
                    current_unit = {
                        'type': parts[0],
                        'index': parts[1],
                        'content': ''
                    }
            elif line and current_unit:
                current_unit['content'] += line + '\n'
        
        # 添加最后一个单元
        if current_unit:
            units.append(current_unit)
        
        return units
    
    def generate_qa_for_units(self, units: List[Dict], output_file: str, delay: float = 1.0):
        """
        为所有单元生成问答对并保存为CSV
        
        Args:
            units: 拆分单元列表
            output_file: 输出CSV文件路径
            delay: API调用间隔（秒）
        """
        all_qa_pairs = []
        
        print(f"开始处理 {len(units)} 个单元...")
        
        for i, unit in enumerate(units):
            print(f"处理单元 {i+1}/{len(units)}: {unit['type']} {unit['index']}")
            
            # 跳过空内容或过短的内容
            content = unit['content'].strip()
            if not content or len(content) < 10:
                print("  跳过：内容过短")
                continue
            
            # 调用LLM生成问答对
            llm_response = self.call_llm(content)
            
            if llm_response:
                # 解析问答对
                qa_pairs = self.parse_qa_pairs(llm_response)
                
                for title, answer, questions in qa_pairs:
                    all_qa_pairs.append({
                        '分段标题': title,
                        '分段内容': answer,
                        '问题': '\n'.join(questions)
                    })
                
                print(f"  生成了 {len(qa_pairs)} 个问答对")
            else:
                print("  生成失败")
            
            # 添加延迟避免API限制
            if i < len(units) - 1:
                time.sleep(delay)
        
        # 保存为CSV
        self.save_to_csv(all_qa_pairs, output_file)
        print(f"\n总共生成了 {len(all_qa_pairs)} 个问答对")
        print(f"结果已保存到: {output_file}")
    
    def save_to_csv(self, qa_pairs: List[Dict], output_file: str):
        """
        保存问答对到CSV文件
        
        Args:
            qa_pairs: 问答对列表
            output_file: 输出文件路径
        """
        with open(output_file, 'w', newline='', encoding='utf-8-sig') as f:
            if qa_pairs:
                writer = csv.DictWriter(f, fieldnames=qa_pairs[0].keys())
                writer.writeheader()
                writer.writerows(qa_pairs)


def main():
    """主函数"""
    # 配置参数
    API_KEY = "sk-j88LiOgtH1b5gYiEut0FgTixSlB9Mh6JdP9qzUlTwaV97ByN"  # 请替换为你的API密钥
    API_BASE = "https://api.openai-proxy.org/v1"  # 可以根据需要修改
    INPUT_FILE = "宪法_拆分结果.txt"
    OUTPUT_FILE = "宪法_RAG问答对.csv"
    
    # 创建生成器
    generator = RAGQAGenerator(API_KEY, API_BASE)
    
    # 读取拆分结果
    print("读取拆分结果文件...")
    units = generator.read_split_result(INPUT_FILE)
    print(f"读取到 {len(units)} 个单元")
    
    # 生成问答对
    print("\n开始生成RAG问答对...")
    generator.generate_qa_for_units(units, OUTPUT_FILE, delay=1.0)
    
    print("\n处理完成！")


if __name__ == "__main__":
    main() 